#include "websocketsession.hxx"

#include "md5.hxx"

#include <boost/bind.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/scoped_array.hpp>

#include <iostream>
#include <vector>
#include <memory>

#include <stdlib.h>

//------------------------------------------------
const char* WEB_SOCKET_DATA_END_FLAG = "\xff";
const char* WEB_SOCKET_HANDSHAKE_END_FLAG = "\r\n\r\n";
const char* HANDSHAKE_HEADER_UPGRADE = "Upgrade"; 
const char* HANDSHAKE_HEADER_CONNECTION = "Connection";
const char* HANDSHAKE_HEADER_HOST = "Host";
const char* HANDSHAKE_HEADER_ORIGIN = "Origin";
const char* HANDSHAKE_HEADER_SEC_WEBSOCKET_KEY1 = "Sec-WebSocket-Key1";
const char* HANDSHAKE_HEADER_SEC_WEBSOCKET_KEY2 = "Sec-WebSocket-Key2";
const char* HANDSHAKE_HEADER_SEC_WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol";

const char* HANDSHAKE_METHOD = "GET";
const char* WEB_SOCKET_SCHEME = "ws://";
const char* HANDSHAKE_UPGRADE_VALUE = "WebSocket";
const char* HANDSHAKE_CONNECTION_VALUE = "Upgrade";

char DATA_START_FLAG = 0x00;
char DATA_END_FLAG = 0xff;

const int CHALLENGE_BYTES = 8;
const int SECKEY_LENGTH = 4;
const int CHALLENGE_RESPONSE_LENGTH = 16;
//------------------------------------------------
class WebSocketSession::HandShakeProcessor
{
private:
   typedef std::vector<std::string> HandShakeInfos;

public:
   HandShakeProcessor(boost::asio::streambuf& headerStream, unsigned int port)
      :mMethod("")
      ,mUri("")
      ,mUpgrade("")
      ,mConnection("")
      ,mHost("")
      ,mOrigin("")
      ,mSecKey1("")
      ,mSeckey2("")
      ,mProtocol("")
      ,mHostPort("")
      ,mPort(port)
   {
      extractHandShakeInfos(headerStream);
   }

public:
   bool validate();
   bool getResponse(std::string& response, unsigned char** finalChallegne);

private:
   void extractHandShakeInfos(boost::asio::streambuf& header);
   bool checkHeaderLines();
   bool checkProtocol();
   bool checkMethod();
   bool checkHost();
   bool generateChallenge();
   bool extractChallenge(std::string& key, unsigned char* secKey);
   void packChallenge(long long digit, unsigned char* fodder);

private:
   std::string mMethod;
   std::string mUri;
   std::string mUpgrade;
   std::string mConnection;
   std::string mHost;
   std::string mOrigin;
   std::string mSecKey1;
   std::string mSeckey2;
   std::string mProtocol;
   std::string mHostPort;

   unsigned char mChallenge[CHALLENGE_BYTES];
   unsigned char mFinalChallenge[CHALLENGE_RESPONSE_LENGTH];
   unsigned int mPort;
};

//private
void WebSocketSession::HandShakeProcessor::extractHandShakeInfos(boost::asio::streambuf& header)
{
   std::istream headerStream(&header);
   std::string line;
   bool isFirstLine = true;
   std::vector<std::string> keyValuePairs;

   while (std::getline(headerStream, line) && line != "\r")
   {
      line.resize(line.length() - 1); //remvoed "\r"
      keyValuePairs.clear();

      if (isFirstLine)
      {
         isFirstLine = false;

         boost::algorithm::split(keyValuePairs, line, boost::is_any_of(" "));
         if (keyValuePairs.size() > 0)
         {
            mMethod = keyValuePairs[0];
            mUri = keyValuePairs[1];
         }
      }
      else
      {
         boost::algorithm::split(keyValuePairs,
            line,
            boost::is_any_of(":"),
            boost::algorithm::token_compress_on);

         if (keyValuePairs.size() > 0)
         {
            if (boost::algorithm::iequals(HANDSHAKE_HEADER_UPGRADE, keyValuePairs[0]))
            {
               mUpgrade = keyValuePairs[1];
               boost::algorithm::trim_left_if(mUpgrade, boost::algorithm::is_any_of(" "));
            }
            else if (boost::algorithm::iequals(HANDSHAKE_HEADER_CONNECTION, keyValuePairs[0]))
            {
               mConnection = keyValuePairs[1];
               boost::algorithm::trim_left_if(mConnection, boost::algorithm::is_any_of(" "));
            }
            else if (boost::algorithm::iequals(HANDSHAKE_HEADER_HOST, keyValuePairs[0]))
            {
               mHostPort = keyValuePairs[2];
               std::string host = keyValuePairs[1];
               boost::algorithm::trim_left_if(host, boost::algorithm::is_any_of(" "));
               mHost += WEB_SOCKET_SCHEME + host + ":" + mHostPort + mUri;
            }
            else if (boost::algorithm::iequals(HANDSHAKE_HEADER_ORIGIN, keyValuePairs[0]))
            {
               mOrigin = keyValuePairs[1];
               boost::algorithm::trim_left_if(mOrigin, boost::algorithm::is_any_of(" "));
            }
            else if (boost::algorithm::iequals(HANDSHAKE_HEADER_SEC_WEBSOCKET_KEY1, keyValuePairs[0]))
            {
               mSecKey1 = keyValuePairs[1];
               boost::algorithm::trim_left_if(mSecKey1, boost::algorithm::is_any_of(" "));
            }
            else if (boost::algorithm::iequals(HANDSHAKE_HEADER_SEC_WEBSOCKET_KEY2, keyValuePairs[0]))
            {
               mSeckey2 = keyValuePairs[1];
               boost::algorithm::trim_left_if(mSeckey2, boost::algorithm::is_any_of(" "));
            }
            else if (boost::algorithm::iequals(HANDSHAKE_HEADER_SEC_WEBSOCKET_PROTOCOL, keyValuePairs[0]))
            {
               mProtocol = keyValuePairs[1];
               boost::algorithm::trim_left_if(mProtocol, boost::algorithm::is_any_of(" "));
            }
         }
      }
   }

   if (!headerStream.eof())
   {
      headerStream.read((char*)mChallenge, CHALLENGE_BYTES);
   }
}

//private
void WebSocketSession::HandShakeProcessor::packChallenge(long long digit, unsigned char* fodder)
{
   fodder[0]=(char)(0xff & (digit >> 24));
   fodder[1]=(char)(0xff & (digit >> 16));
   fodder[2]=(char)(0xff & (digit >> 8));
   fodder[3]=(char)(0xff & digit);

   //std::cout << fodder <<std::endl;
}

//private
bool WebSocketSession::HandShakeProcessor::extractChallenge(std::string& key, unsigned char* secKey)
{
   bool rt = false;

   int spaceCount = 0;
   long long digital = 0;

   for (int i = 0; i < key.length(); ++i)
   {
      if (boost::algorithm::is_digit()(key[i]))
      {
         digital = digital*10 + (key[i] - '0');
      }

      if (boost::algorithm::is_space()(key[i]))
      {
         ++spaceCount;
      }
   }

   if (spaceCount > 0)
   {
      if (digital % spaceCount == 0)
      {
         packChallenge(digital / spaceCount, secKey);

         rt = true;
      }      else      {         std::cout << __FUNCTION__ << " | key is not an integral"            << " | key: " << digital            << " | space: " << spaceCount << std::endl;       }
   }

   return rt;
}

//private
bool WebSocketSession::HandShakeProcessor::generateChallenge()
{
   bool rt = false;

   if (mSecKey1.length() > 0 && mSeckey2.length() > 0)
   {
      unsigned char secKey1Tmp[SECKEY_LENGTH];
      unsigned char secKey2Tmp[SECKEY_LENGTH];

      if (extractChallenge(mSecKey1, secKey1Tmp) &&
         extractChallenge(mSeckey2, secKey2Tmp))
      {
         unsigned char tmp[CHALLENGE_RESPONSE_LENGTH];
         memcpy(tmp, secKey1Tmp, SECKEY_LENGTH);
         memcpy(tmp + SECKEY_LENGTH, secKey2Tmp, SECKEY_LENGTH);
         memcpy(tmp + CHALLENGE_BYTES, mChallenge, CHALLENGE_BYTES);

         MD5 md5;
         md5.update(tmp, CHALLENGE_RESPONSE_LENGTH);
         md5.finalize();
         unsigned char* md = md5.raw_digest();
         if (NULL != md)
         {
            memcpy(mFinalChallenge, md, CHALLENGE_RESPONSE_LENGTH);
            delete[] md;

            rt = true;
         }
      }
   }
   else
   {
      rt = true;
   }

   return rt;
}

//private
bool WebSocketSession::HandShakeProcessor::checkMethod()
{
   bool rt = boost::algorithm::iequals(HANDSHAKE_METHOD, mMethod);

   if (!rt)
   {
      std::cout << __FUNCTION__ << " | not support method now: " << mMethod << std::endl;
   }

   return rt;
}

//private
bool WebSocketSession::HandShakeProcessor::checkHeaderLines()
{
   bool rt = boost::algorithm::iequals(mConnection, HANDSHAKE_CONNECTION_VALUE) &&
      boost::algorithm::iequals(mUpgrade, HANDSHAKE_UPGRADE_VALUE);

   if (!rt)
   {
      std::cout << __FUNCTION__ << " | connection or upgrade Header is not defined" << std::endl;
   }

   return rt;
}

//private
bool WebSocketSession::HandShakeProcessor::checkProtocol()
{
   bool rt = true;

   if (mProtocol.length() > 0)
   {
      for (int i = 0; i < mProtocol.length(); ++i)
      {
         if ((int)mProtocol[i] < 32 || (int)mProtocol[i] > 126)
         {
            std::cout << __FUNCTION__ << " | invalid protocol: " << mProtocol << std::endl;

            rt = false;
            break;
         }
      }
   }

   return rt;
}

//private
bool WebSocketSession::HandShakeProcessor::checkHost()
{
   bool rt = (::atol(mHostPort.c_str()) == mPort);

   if (!rt)
   {
      std::cout << __FUNCTION__ << " | invalid port: " << mHostPort << std::endl;
   }

   return rt;
}

//public
bool WebSocketSession::HandShakeProcessor::validate()
{
   return (checkMethod() &&
      checkHeaderLines() &&
      checkProtocol() &&
      checkHost());
}

//public
bool WebSocketSession::HandShakeProcessor::getResponse(std::string& response, unsigned char** finalChallegne)
{
   bool rt = false;

   if (generateChallenge())
   {
      response += "HTTP/1.1 101 WebSocket Protocol Handshake\r\n";
      response += "Upgrade: WebSocket\r\n";
      response += "Connection: Upgrade\r\n";

      response += "Sec-WebSocket-Location: " + mHost + "\r\n";
      response += "Sec-WebSocket-Origin: " + mOrigin + "\r\n";

      if (mProtocol.length() > 0)
      {
         response += "Sec-WebSocket-Protocol: " + mProtocol + "\r\n";
      }

      response += "\r\n";

      *finalChallegne = mFinalChallenge;

      rt = true;
   }

   return rt;
}
//------------------------------------------------
WebSocketSession::WebSocketSession(boost::asio::io_service& ioService, unsigned int port)
:mSocket(ioService)
,mStrand(ioService)
,mHandShakeResponse("")
,mPort(port)
,mCurrenStatus(S_START)
,mDataTransferType(DTT_UNKNOWN)
{
}

WebSocketSession::~WebSocketSession(void)
{
}
//------------------------------------------------
//public
void WebSocketSession::start()
{
   recvHandShakeRequest();
}
//------------------------------------------------
//private
void WebSocketSession::handleRead(const boost::system::error_code& error, size_t bytesTransferred)
{
   std::cout << __FUNCTION__ << " | bytes:" << bytesTransferred << std::endl;

   if (0 == error)
   {
      switch(mCurrenStatus)
      {
      case S_START:
         {
            handleHandShakeResponse();
         }
         break;
      case S_HANDSHAKE_SUCCESSED:
         {
            setDataTransferType();
         }
         break;
      case S_DATA_TRANSFER:
         {
            mFinalData.reset(new unsigned char[bytesTransferred + 2]);
            if (mFinalData)
            {
               memcpy(mFinalData.get(), &DATA_START_FLAG, 1);
               std::istream dataStream(&mData);
               dataStream.read((char*)mFinalData.get() + 1, bytesTransferred);
               memcpy(mFinalData.get() + bytesTransferred + 1, &DATA_END_FLAG, 1);

               std::cout << __FUNCTION__ << " | data transfer: " << mFinalData.get() << std::endl;

               boost::asio::async_write(mSocket,
                  boost::asio::buffer(mFinalData.get(), bytesTransferred + 2),
                  mStrand.wrap(boost::bind(&WebSocketSession::handleWrite,
                  shared_from_this(),
                  boost::asio::placeholders::error
                  )));
            }
         }
         break;
      }
   }
   else
   {
      std::cout << __FUNCTION__ <<  " | handleRead error! error code: " <<
         error.value() << " | error msg: " << error.message() << std::endl;

      mSocket.close();
   }
}

//private
void WebSocketSession::handleWrite(const boost::system::error_code& error)
{
   if ( 0 == error)
   {
      switch(mCurrenStatus)
      {
      case S_HANDSHAKE_HEADER:
         {
            mCurrenStatus = S_HANDSHAKE_SUCCESSED;

            recvDataType();
         }
         break;
      case S_DATA_TRANSFER:
         {
            boost::asio::async_read_until(mSocket,
               mData,
               WEB_SOCKET_DATA_END_FLAG,
               mStrand.wrap(boost::bind(&WebSocketSession::handleRead,
               shared_from_this(),
               boost::asio::placeholders::error,
               boost::asio::placeholders::bytes_transferred
               )));
         }
         break;
      }
   }
   else
   {
      std::cout << "handWrite error! error code: " << error.value()
         << " | error msg: " << error.message() << std::endl;
      mSocket.close();
   }
}

//private
void WebSocketSession::setDataTransferType()
{
   std::cout << __FUNCTION__ << " | data transfer type: " << (int)mDataType[0] << std::endl;

   mCurrenStatus = S_DATA_TRANSFER;
   if ((int)mDataType[0] & 0x80)
   {
      mDataTransferType = DTT_BINARY;
   }
   else if ((int)mDataType[0] == 0x00)
   {
      mDataTransferType = DTT_TEXT;

      boost::asio::async_read_until(mSocket,
         mData,
         WEB_SOCKET_DATA_END_FLAG,
         mStrand.wrap(boost::bind(&WebSocketSession::handleRead,
         shared_from_this(),
         boost::asio::placeholders::error,
         boost::asio::placeholders::bytes_transferred
         )));
   }
}

//private
void WebSocketSession::recvDataType()
{
   mSocket.async_read_some(boost::asio::buffer(mDataType, DTATA_TRANSFER_TYPE_LENGTH),
      mStrand.wrap(boost::bind(&WebSocketSession::handleRead,
      shared_from_this(),
      boost::asio::placeholders::error,
      boost::asio::placeholders::bytes_transferred
      )));
}

//private
void WebSocketSession::recvHandShakeRequest()
{
   boost::asio::async_read_until(mSocket,
      mHeader,
      WEB_SOCKET_HANDSHAKE_END_FLAG,
      mStrand.wrap(boost::bind(&WebSocketSession::handleRead,
      shared_from_this(),
      boost::asio::placeholders::error,
      boost::asio::placeholders::bytes_transferred
      )));
}

//private
void WebSocketSession::handleHandShakeResponse()
{
   boost::scoped_ptr<HandShakeProcessor> handShakeProcessor(new HandShakeProcessor(mHeader, mPort));
   unsigned char* finalChallenge;
   if (handShakeProcessor &&
      handShakeProcessor->validate() &&
      handShakeProcessor->getResponse(mHandShakeResponse, &finalChallenge))
   {
      mCurrenStatus = S_HANDSHAKE_HEADER;

      memcpy(mFinalHandShakeResponse, mHandShakeResponse.c_str(), mHandShakeResponse.length());
      memcpy(mFinalHandShakeResponse + mHandShakeResponse.length(), finalChallenge, CHALLENGE_RESPONSE_LENGTH);

      std::cout << __FUNCTION__ << " | mHandShakeResponse length: " << mHandShakeResponse.length() + CHALLENGE_RESPONSE_LENGTH 
         << " | data: \r\n" << mFinalHandShakeResponse << std::endl;

      boost::asio::async_write(mSocket,
         boost::asio::buffer(mFinalHandShakeResponse, mHandShakeResponse.length() + CHALLENGE_RESPONSE_LENGTH),
         mStrand.wrap(boost::bind(&WebSocketSession::handleWrite,
         shared_from_this(),
         boost::asio::placeholders::error
         )));
   }
   else
   {
      std::cout << __FUNCTION__ << " | process handshake failed" << std::endl;
      mSocket.close();
   }
}
//------------------------------------------------